extern crate base64;
use sha1::Digest;
use std::borrow::Cow;

use aes::Aes256;
use aes::cipher::generic_array::GenericArray;
use aes::cipher::{BlockDecrypt, BlockEncryptMut, KeyInit, KeyIvInit};
use anyhow::{Result, anyhow};
use base64::Engine;
use base64::engine::general_purpose;
use cbc::cipher::block_padding::NoPadding;
use rand::{Rng, rng};
use sha1::Sha1;

pub fn base64_decode(cypher: &str) -> Result<Vec<u8>> {
    general_purpose::STANDARD
        .decode(cypher)
        .map_err(|e| anyhow!(e))
}

pub fn parse_decrypted(decrypted_bytes: &[u8]) -> Result<(Cow<str>, usize, Cow<str>, Cow<str>)> {
    let random = String::from_utf8_lossy(&decrypted_bytes[0..16]);

    let msg_len = u32::from_be_bytes([
        decrypted_bytes[16],
        decrypted_bytes[17],
        decrypted_bytes[18],
        decrypted_bytes[19],
    ]) as usize;

    let msg = String::from_utf8_lossy(&decrypted_bytes[20..20 + msg_len]);
    let app_id = String::from_utf8_lossy(&decrypted_bytes[20 + msg_len..]);
    Ok((random, msg_len, msg, app_id))
}

const BASE: &str = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";

pub fn get_random_str() -> String {
    let mut rng = rng();
    let base_chars: Vec<char> = BASE.chars().collect();

    (0..16)
        .map(|_| {
            let index = rng.random_range(0..base_chars.len());
            base_chars[index]
        })
        .collect()
}

// 生成4个字节的网络字节序
fn get_network_bytes_order(source_number: usize) -> [u8; 4] {
    [
        (source_number >> 24) as u8,
        (source_number >> 16) as u8,
        (source_number >> 8) as u8,
        source_number as u8,
    ]
}

// PKCS7填充
fn pkcs7_encode(data_len: usize) -> Vec<u8> {
    let block_size = 32; // AES-256-CBC 块大小为32字节
    let pad_len = block_size - (data_len % block_size);
    vec![pad_len as u8; pad_len]
}

type Aes256Cbc = cbc::Encryptor<Aes256>;

pub fn do_encrypt(random_str: &str, reply_msg: &str, app_id: &str, encoding_aes_key: &str) -> String {
    // 准备各个组件的字节数组
    let random_str_bytes = random_str.as_bytes();
    let text_bytes = reply_msg.as_bytes();
    let network_bytes_order = get_network_bytes_order(text_bytes.len());
    let app_id_bytes = app_id.as_bytes();

    // 拼接字节: randomStr + networkBytesOrder + text + appid
    let mut byte_collector = Vec::new();
    byte_collector.extend_from_slice(random_str_bytes);
    byte_collector.extend_from_slice(&network_bytes_order);
    byte_collector.extend_from_slice(text_bytes);
    byte_collector.extend_from_slice(app_id_bytes);

    // 添加PKCS7填充
    let pad_bytes = pkcs7_encode(byte_collector.len());
    byte_collector.extend_from_slice(&pad_bytes);

    let aes_key_bytes =
        base64_decode(encoding_aes_key).expect("encoding_aes_key base64 decode error !!!");
    if aes_key_bytes.len() != 32 {
        panic!("AES key must be 32 bytes long");
    }

    // 使用前16字节作为IV
    let iv = &aes_key_bytes[0..16];

    // 执行AES-256-CBC加密
    let mut buffer = byte_collector.clone();
    let cipher = Aes256Cbc::new_from_slices(&aes_key_bytes, iv)
        .expect("AES encryption algorithm is incorrect");
    let encrypted = cipher
        .encrypt_padded_mut::<NoPadding>(&mut buffer, byte_collector.len())
        .expect("AES encryption algorithm is incorrect");

    // Base64编码
    general_purpose::STANDARD.encode(&encrypted)
}

#[test]
pub fn test_encrypt() {
    let random_str = "707722b803182950";
    let msg = "{\"demo_resp\":\"good luck\"}";
    let app_id = "wxba5fad812f8e6fb9";
    let encoding_aes_key = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
    let string = do_encrypt(random_str, msg, app_id, encoding_aes_key);
    assert_eq!(
        "ELGduP2YcVatjqIS+eZbp80MNLoAUWvzzyJxgGzxZO/5sAvd070Bs6qrLARC9nVHm48Y4hyRbtzve1L32tmxSQ==",
        string
    );
}

#[test]
pub fn test_decrypt() {
    let encoding_aes_key = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
    let encrypted_msg = "+qdx1OKCy+5JPCBFWw70tm0fJGb2Jmeia4FCB7kao+/Q5c/ohsOzQHi8khUOb05JCpj0JB4RvQMkUyus8TPxLKJGQqcvZqzDpVzazhZv6JsXUnnR8XGT740XgXZUXQ7vJVnAG+tE8NUd4yFyjPy7GgiaviNrlCTj+l5kdfMuFUPpRSrfMZuMcp3Fn2Pede2IuQrKEYwKSqFIZoNqJ4M8EajAsjLY2km32IIjdf8YL/P50F7mStwntrA2cPDrM1kb6mOcfBgRtWygb3VIYnSeOBrebufAlr7F9mFUPAJGj04=";
    let encrypted_bytes = base64_decode(encrypted_msg)
        .expect("encrypt_msg base64 decode error !!!");
    let aes_bytes = base64_decode(encoding_aes_key).expect("encoding_aes_key base64 decode error !!!");
    let vec = do_decrypt(&encrypted_bytes, &aes_bytes);
    let decrypted = parse_decrypted(&vec).expect("Should be able to decrypt");

    assert_eq!(decrypted.0, "a8eedb185eb2fecf");
    assert_eq!(decrypted.1, 167);
    assert_eq!(decrypted.2, "{\"ToUserName\":\"gh_97417a04a28d\",\"FromUserName\":\"o9AgO5Kd5ggOC-bXrbNODIiE3bGY\",\"CreateTime\":1714112445,\"MsgType\":\"event\",\"Event\":\"debug_demo\",\"debug_str\":\"hello world\"}");
    assert_eq!(decrypted.3, "wxba5fad812f8e6fb9");
}

pub fn do_decrypt(encrypted_bytes: &[u8], secret_key_bytes: &[u8]) -> Vec<u8> {
    let iv = secret_key_bytes[..16].to_vec();

    let cipher = Aes256::new(secret_key_bytes.into());
    let mut decrypted_blocks: Vec<Vec<u8>> = Vec::new();
    (0..encrypted_bytes.len()).step_by(16).for_each(|x| {
        let last = if x == 0 {
            &iv
        } else {
            &encrypted_bytes[x - 16..x]
        };

        let mut block = GenericArray::clone_from_slice(&encrypted_bytes[x..x + 16]);
        cipher.decrypt_block(&mut block);
        let decrypted_block: Vec<u8> = block.into_iter().collect();
        let xor_block = xor_bytes(last, &decrypted_block);
        decrypted_blocks.push(xor_block);
    });

    // Get number of padding bytes applied during encryption & remove padding
    let padding_byte = *decrypted_blocks
        .last()
        .and_then(|block| block.last())
        .unwrap() as usize;

    decrypted_blocks
        .into_iter()
        .flatten()
        .take(encrypted_bytes.len() - padding_byte)
        .collect()
}

fn xor_bytes(bytes1: &[u8], bytes2: &[u8]) -> Vec<u8> {
    bytes1
        .iter()
        .zip(bytes2.iter())
        .map(|(&b1, &b2)| b1 ^ b2)
        .collect()
}

// 验证消息签名
pub fn verify_msg_signature(seeds: Vec<&str>, expected_signature: &str) -> bool {
    let sign = make_sign(seeds);
    tracing::info!("calc_sign: {:?}", sign);
    sign == expected_signature
}

pub fn make_sign(mut seeds: Vec<&str>) -> String {
    seeds.sort_unstable();

    let concatenated = seeds.concat();
    let mut hasher = Sha1::new();
    hasher.update(concatenated.as_bytes());
    let result = hasher.finalize();
    hex::encode(result)
}
